# -*- coding: utf-8 -*-
"""
图像处理工具模块
"""

import base64
import io
from PIL import Image, ImageDraw
import numpy as np
import asyncio
from config import is_red, is_lvl, width, height, use_color_mode

# 定义调色板数组
pal_arr = [
    [[0, 0, 0], [255, 255, 255]],
    [[0, 0, 0], [255, 255, 255], [127, 0, 0]],
    [[0, 0, 0], [255, 255, 255], [127, 127, 127]],
    [[0, 0, 0], [255, 255, 255], [127, 127, 127], [127, 0, 0]],
    [[0, 0, 0], [255, 255, 255]],
    [[0, 0, 0], [255, 255, 255], [220, 180, 0]],
    [[0, 0, 0]],
    [[0, 0, 0], [255, 255, 255], [0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 0], [255, 128, 0]]
]


def get_near(r, g, b, cur_pal):
    """
    获取最接近的颜色索引
    :param r: 红色分量
    :param g: 绿色分量
    :param b: 蓝色分量
    :param cur_pal: 当前调色板
    :return: 最接近的颜色索引
    """
    ind = 0
    err = get_err(r, g, b, cur_pal[0])
    for i in range(1, len(cur_pal)):
        cur = get_err(r, g, b, cur_pal[i])
        if cur < err:
            err = cur
            ind = i
    return ind


def get_err(r, g, b, std_col):
    """
    计算颜色误差
    :param r: 红色分量
    :param g: 绿色分量
    :param b: 蓝色分量
    :param std_col: 标准颜色
    :return: 误差值
    """
    r -= std_col[0]
    g -= std_col[1]
    b -= std_col[2]
    return r * r + g * g + b * b


def add_val(c, r, g, b, k):
    """
    添加误差扩散值
    :param c: 颜色数组
    :param r: 红色误差
    :param g: 绿色误差
    :param b: 蓝色误差
    :param k: 系数
    :return: 新的颜色数组
    """
    return [c[0] + (r * k) / 32, c[1] + (g * k) / 32, c[2] + (b * k) / 32]


def set_val(p_data, i, c, cur_pal):
    """
    设置像素值
    :param p_data: 像素数据
    :param i: 索引
    :param c: 颜色索引
    :param cur_pal: 当前调色板
    """
    p_data[i] = cur_pal[c][0]
    p_data[i + 1] = cur_pal[c][1]
    p_data[i + 2] = cur_pal[c][2]
    p_data[i + 3] = 255


def process_image(image_base64):
    """
    处理图片 - 灰度化或抖动

    :param image_base64: 图片的base64字符串
    :return: 处理后的图片base64字符串
    """
    try:
        # 解码base64图片
        image_data = base64.b64decode(image_base64.split(',')[1] if ',' in image_base64 else image_base64)
        src_img = Image.open(io.BytesIO(image_data))
        
        # 获取图像尺寸
        s_w, s_h = src_img.size
        
        # 设置边界参数
        d_x, d_y = 0, 0
        d_w, d_h = s_w, s_h
        
        # 创建目标图像
        dst_img = Image.new('RGBA', (d_w, d_h), (0, 0, 0, 0))
        
        # 选择调色板
        pal_ind = 1 if is_red else 0
        cur_pal = pal_arr[pal_ind]
        
        # 获取源图像数据
        src_pixels = np.array(src_img.convert('RGBA'))
        dst_pixels = np.zeros((d_h, d_w, 4), dtype=np.uint8)
        
        if is_lvl:
            # 灰度处理
            for j in range(d_h):
                y = d_y + j
                if y < 0 or y >= s_h:
                    for i in range(d_w):
                        idx = (j * d_w + i) * 4
                        c = 1 if (i + j) % 2 == 0 else 0
                        dst_pixels[j, i, 0] = cur_pal[c][0]
                        dst_pixels[j, i, 1] = cur_pal[c][1]
                        dst_pixels[j, i, 2] = cur_pal[c][2]
                        dst_pixels[j, i, 3] = 255
                    continue
                
                for i in range(d_w):
                    x = d_x + i
                    if x < 0 or x >= s_w:
                        c = 1 if (i + j) % 2 == 0 else 0
                        dst_pixels[j, i, 0] = cur_pal[c][0]
                        dst_pixels[j, i, 1] = cur_pal[c][1]
                        dst_pixels[j, i, 2] = cur_pal[c][2]
                        dst_pixels[j, i, 3] = 255
                        continue
                    
                    r, g, b = src_pixels[y, x, 0], src_pixels[y, x, 1], src_pixels[y, x, 2]
                    c = get_near(r, g, b, cur_pal)
                    dst_pixels[j, i, 0] = cur_pal[c][0]
                    dst_pixels[j, i, 1] = cur_pal[c][1]
                    dst_pixels[j, i, 2] = cur_pal[c][2]
                    dst_pixels[j, i, 3] = 255
        else:
            # 抖动处理（Floyd-Steinberg算法）
            a_ind, b_ind = 0, 1
            err_arr = np.zeros((2, d_w, 3))
            
            for j in range(d_h):
                y = d_y + j
                if y < 0 or y >= s_h:
                    for i in range(d_w):
                        c = 1 if (i + j) % 2 == 0 else 0
                        dst_pixels[j, i, 0] = cur_pal[c][0]
                        dst_pixels[j, i, 1] = cur_pal[c][1]
                        dst_pixels[j, i, 2] = cur_pal[c][2]
                        dst_pixels[j, i, 3] = 255
                    continue
                
                a_ind, b_ind = b_ind, a_ind
                err_arr[b_ind] = np.zeros((d_w, 3))
                
                for i in range(d_w):
                    x = d_x + i
                    if x < 0 or x >= s_w:
                        c = 1 if (i + j) % 2 == 0 else 0
                        dst_pixels[j, i, 0] = cur_pal[c][0]
                        dst_pixels[j, i, 1] = cur_pal[c][1]
                        dst_pixels[j, i, 2] = cur_pal[c][2]
                        dst_pixels[j, i, 3] = 255
                        continue
                    
                    old = err_arr[a_ind, i]
                    r = float(src_pixels[y, x, 0]) + old[0]
                    g = float(src_pixels[y, x, 1]) + old[1]
                    b = float(src_pixels[y, x, 2]) + old[2]
                    
                    col_ind = get_near(r, g, b, cur_pal)
                    col_val = cur_pal[col_ind]
                    
                    dst_pixels[j, i, 0] = col_val[0]
                    dst_pixels[j, i, 1] = col_val[1]
                    dst_pixels[j, i, 2] = col_val[2]
                    dst_pixels[j, i, 3] = 255
                    
                    r_err = r - col_val[0]
                    g_err = g - col_val[1]
                    b_err = b - col_val[2]
                    
                    # 误差扩散
                    if i == 0:
                        err_arr[b_ind, i] += np.array([r_err, g_err, b_err]) * 7.0 / 32
                        err_arr[b_ind, i + 1] += np.array([r_err, g_err, b_err]) * 2.0 / 32
                        err_arr[a_ind, i + 1] += np.array([r_err, g_err, b_err]) * 7.0 / 32
                    elif i == d_w - 1:
                        err_arr[b_ind, i - 1] += np.array([r_err, g_err, b_err]) * 7.0 / 32
                        err_arr[b_ind, i] += np.array([r_err, g_err, b_err]) * 9.0 / 32
                    else:
                        err_arr[b_ind, i - 1] += np.array([r_err, g_err, b_err]) * 3.0 / 32
                        err_arr[b_ind, i] += np.array([r_err, g_err, b_err]) * 5.0 / 32
                        err_arr[b_ind, i + 1] += np.array([r_err, g_err, b_err]) * 1.0 / 32
                        err_arr[a_ind, i + 1] += np.array([r_err, g_err, b_err]) * 7.0 / 32
        
        # 创建处理后的图像
        result_img = Image.fromarray(dst_pixels)
        
        # 转换为base64并返回
        buffered = io.BytesIO()
        result_img.save(buffered, format="PNG")
        result_base64 = f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode('utf-8')}"
        
        return result_base64
    except Exception as e:
        raise Exception(f"图片处理失败: {str(e)}")


def _process_base64_to_pixel_array(base64_str):
    """
    内部函数：将base64图片字符串转换为像素数组
    :param base64_str: base64格式的图片字符串
    :return: 包含像素值的数组
    """
    # 解码base64图片
    image_data = base64.b64decode(base64_str.split(',')[1] if ',' in base64_str else base64_str)
    img = Image.open(io.BytesIO(image_data))
    
    # 调整图片大小
    img = img.resize((width, height))
    
    # 转换为RGBA模式
    img = img.convert('RGBA')
    
    # 获取图像数据
    img_data = np.array(img)
    pixel_array = [0] * (width * height)
    
    # 处理图像数据
    for y in range(height):
        for x in range(width):
            i = y * width + x
            if use_color_mode:
                # 7色模式
                pixel_array[i] = get_val_7color(img_data[y, x])
            else:
                # 黑白灰模式
                pixel_array[i] = get_val(img_data[y, x])
    
    return pixel_array


def base64_to_pixel_array(base64_str):
    """
    同步将base64图片字符串转换为像素数组
    :param base64_str: base64格式的图片字符串
    :return: 包含像素值的数组
    """
    try:
        return _process_base64_to_pixel_array(base64_str)
    except Exception as e:
        raise Exception(f"图片转换为像素数组失败: {str(e)}")


async def base64_to_pixel_array_async(base64_str):
    """
    异步将base64图片字符串转换为像素数组
    :param base64_str: base64格式的图片字符串
    :return: 包含像素值的数组
    """
    try:
        # 使用线程池执行CPU密集型操作
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            None, _process_base64_to_pixel_array, base64_str
        )
    except Exception as e:
        raise Exception(f"图片转换为像素数组失败: {str(e)}")


def get_val(pixel):
    """
    获取像素值（黑白灰模式）
    :param pixel: 像素数据
    :return: 返回像素值：0(黑)、1(白)、2(灰)或3(其他)
    """
    r, g, b = pixel[0], pixel[1], pixel[2]
    if r == 0x00 and g == 0x00:
        return 0
    if r == 0xFF and g == 0xFF:
        return 1
    if r == 0x7F and g == 0x7F:
        return 2
    return 3


def get_val_7color(pixel):
    """
    获取7色模式下的像素值
    :param pixel: 像素数据
    :return: 返回像素值：0(黑)、1(白)、2(绿)、3(蓝)、4(红)、5(黄)、6(橙)或7(其他)
    """
    r, g, b = pixel[0], pixel[1], pixel[2]
    if r == 0x00 and g == 0x00 and b == 0x00:
        return 0
    if r == 0xFF and g == 0xFF and b == 0xFF:
        return 1
    if r == 0x00 and g == 0xFF and b == 0x00:
        return 2
    if r == 0x00 and g == 0x00 and b == 0xFF:
        return 3
    if r == 0xFF and g == 0x00 and b == 0x00:
        return 4
    if r == 0xFF and g == 0xFF and b == 0x00:
        return 5
    if r == 0xFF and g == 0x80 and b == 0x00:
        return 6
    return 7